import numpy as np
import pygltflib as gltf2

from vtkmodules.vtkCommonMath  import vtkQuaternionf

from vtkmodules.vtkCommonDataModel import (
    vtkPolyData,
    vtkCellArray,
)

from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper
)

from vtkmodules.vtkCommonCore import (
    vtkPoints,
)

# get gltf accessor data as numpy array
def get_accessor_data(gltf, index):
    if gltf is None:
        return None
    accessor = gltf.accessors[index]
    # component basic data type and size
    dtype, bytes_of_component = np.byte, 1
    if accessor.componentType == gltf2.BYTE:
        dtype, bytes_of_component  = np.byte, 1
    if accessor.componentType == gltf2.UNSIGNED_BYTE:        
        dtype, bytes_of_component  = np.uint8, 1
    if accessor.componentType == gltf2.SHORT:        
        dtype, bytes_of_component  = np.short, 2
    if accessor.componentType == gltf2.UNSIGNED_SHORT:        
        dtype, bytes_of_component  = np.ushort, 2
    if accessor.componentType == gltf2.UNSIGNED_INT:        
        dtype, bytes_of_component  = np.uint32, 4
    if accessor.componentType == gltf2.FLOAT:        
        dtype, bytes_of_component  = np.float32, 4

    num_of_components = 1
    if accessor.type == gltf2.SCALAR:
        num_of_components = 1
    if accessor.type == gltf2.VEC2:
        num_of_components = 2
    if accessor.type == gltf2.VEC3:
        num_of_components = 3
    if accessor.type == gltf2.VEC4:
        num_of_components = 4
    if accessor.type == gltf2.MAT2:
        num_of_components = 4
    if accessor.type == gltf2.MAT3:
        num_of_components = 9
    if accessor.type == gltf2.MAT4:
        num_of_components = 16    

    bufferView = gltf.bufferViews[accessor.bufferView]
    bufferData = gltf.get_data_from_buffer_uri(gltf.buffers[bufferView.buffer].uri)    

    # start offset
    offset = bufferView.byteOffset + accessor.byteOffset
    
    size = num_of_components
    if bufferView.byteStride:
        # TODO: check if it is OK to cast strided data...
        size = int(max(bufferView.byteStride, num_of_components * bytes_of_component) / bytes_of_component)

    # use numpy
    data = np.frombuffer(bufferData, dtype, accessor.count * size, offset)        
    
    # reshape 
    data = data.reshape([accessor.count, size])
    # stride
    data = data[:,:num_of_components]

    if accessor.type == gltf2.SCALAR:
        return data[:,0]    

    if accessor.type == gltf2.MAT2:
        data = data.reshape([accessor.count, 2, 2])   
    if accessor.type == gltf2.MAT3:
        data = data.reshape([accessor.count, 3, 3])   
    if accessor.type == gltf2.MAT4:
        data = data.reshape([accessor.count, 4, 4])    
    
    return data


def get_rotation_mat(rotation):
    quat = vtkQuaternionf(rotation[3], rotation[0], rotation[1], rotation[2])
    r = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    quat.ToMatrix3x3(r)
    mat = [
        r[0][0], r[0][1], r[0][2], 0,
        r[1][0], r[1][1], r[1][2], 0,
        r[2][0], r[2][1], r[2][2], 0,
        0, 0, 0, 1]
    return mat

def get_rotation_mat_interp(r0, r1, f):
    q0 = vtkQuaternionf(r0[3], r0[0], r0[1], r0[2])
    q1 = vtkQuaternionf(r1[3], r1[0], r1[1], r1[2])
    q = q0.Slerp(f, q1)
    r = [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
    q.ToMatrix3x3(r)
    mat = [
        r[0][0], r[0][1], r[0][2], 0,
        r[1][0], r[1][1], r[1][2], 0,
        r[2][0], r[2][1], r[2][2], 0,
        0, 0, 0, 1]
    return mat



class BBoxActor(vtkActor):
    def __init__(self, actor) -> None:
        super().__init__()
        self.actor = actor
        
        self.points   = vtkPoints()
        self.points.SetNumberOfPoints(8)

        cells = vtkCellArray()
        cells.InsertNextCell(2, [0, 1])
        cells.InsertNextCell(2, [0, 2])
        cells.InsertNextCell(2, [3, 1])
        cells.InsertNextCell(2, [3, 2])
        cells.InsertNextCell(2, [4, 5])
        cells.InsertNextCell(2, [4, 6])
        cells.InsertNextCell(2, [7, 5])
        cells.InsertNextCell(2, [7, 6])
        cells.InsertNextCell(2, [0, 4])
        cells.InsertNextCell(2, [1, 5])
        cells.InsertNextCell(2, [2, 6])
        cells.InsertNextCell(2, [3, 7])

        self.polyData = vtkPolyData()
        self.polyData.SetPoints(self.points)
        self.polyData.SetLines(cells)

        self.mapper = vtkPolyDataMapper()
        self.mapper.SetInputData(self.polyData)
        self.SetMapper(self.mapper)

        self.Update()

    def Update(self):
        bounds = self.actor.GetBounds()
        x = bounds[0:2]
        y = bounds[2:4]
        z = bounds[4:6]
        self.points.SetPoint(0, [x[0], y[0], z[0]])
        self.points.SetPoint(1, [x[0], y[0], z[1]])
        self.points.SetPoint(2, [x[0], y[1], z[0]])
        self.points.SetPoint(3, [x[0], y[1], z[1]])
        self.points.SetPoint(4, [x[1], y[0], z[0]])
        self.points.SetPoint(5, [x[1], y[0], z[1]])
        self.points.SetPoint(6, [x[1], y[1], z[0]])
        self.points.SetPoint(7, [x[1], y[1], z[1]])
        self.points.Modified()